/*
 * UBI manager for safely updating static volumes (rootfs/assets)
 */

#include "ubi_manager.h"
#include "utils_log.h"

#include <sys/mount.h>
#include <unistd.h>
#include <fcntl.h>
#include <cstdio>
#include <cerrno>
#include <cstring>
#include <fstream>
#include <string>

#include <libubi.h>

namespace El {
namespace Upgrade {

namespace {
constexpr const char *kUbiCtrl = "/dev/ubi_ctrl";
}

UbiManager::UbiManager() {}
UbiManager::~UbiManager() {}

bool UbiManager::IsMountPointMounted(const std::string &mount_point)
{
    std::ifstream mounts("/proc/mounts");
    if (!mounts.is_open()) return false;
    std::string spec, mp, fstype, rest;
    while (mounts) {
        spec.clear(); mp.clear(); fstype.clear(); rest.clear();
        mounts >> spec >> mp >> fstype;
        std::getline(mounts, rest);
        if (mp == mount_point) return true;
    }
    return false;
}

bool UbiManager::UmountIfMounted(const std::string &mount_point)
{
    if (!IsMountPointMounted(mount_point)) {
        return true;
    }
    EL_INFO("Unmounting {} before upgrade", mount_point);
    if (umount2(mount_point.c_str(), MNT_FORCE) == 0) {
        return true;
    }
    EL_ERROR("Failed to unmount {}: {}", mount_point, strerror(errno));
    return false;
}

bool UbiManager::FindUbiNumByMtd(int mtd_num, int &ubi_num)
{
    libubi_t ubi = libubi_open();
    if (!ubi) {
        EL_WARN("libubi_open failed (is UBI enabled?)");
        return false;
    }
    int ret = mtd_num2ubi_dev(ubi, mtd_num, &ubi_num);
    libubi_close(ubi);
    return (ret == 0);
}

bool UbiManager::RemoveUbiblockIfExists(int ubi_num, int vol_id)
{
    char vol_node[64];
    snprintf(vol_node, sizeof(vol_node), "/dev/ubi%d_%d", ubi_num, vol_id);
    int fd = open(vol_node, O_RDONLY);
    if (fd < 0) {
        // volume node not present; nothing to remove
        EL_INFO("Volume node {} not openable, skip ubiblock removal", vol_node);
        return true;
    }
    int rc = ubi_vol_block_remove(fd);
    if (rc != 0) {
        EL_WARN("ubi_vol_block_remove failed for {}: {}", vol_node, strerror(errno));
        close(fd);
        // Not fatal if block not present
        return true;
    }
    close(fd);
    return true;
}

bool UbiManager::DetachMtd(int mtd_num)
{
    libubi_t ubi = libubi_open();
    if (!ubi) {
        EL_ERROR("libubi_open failed in DetachMtd");
        return false;
    }
    int rc = ubi_detach_mtd(ubi, kUbiCtrl, mtd_num);
    if (rc != 0) {
        EL_WARN("ubi_detach_mtd(mtd{}): {} (may already be detached)", mtd_num, strerror(errno));
        libubi_close(ubi);
        return true; // tolerate
    }
    libubi_close(ubi);
    return true;
}

bool UbiManager::AttachMtd(int mtd_num, int &ubi_num)
{
    libubi_t ubi = libubi_open();
    if (!ubi) {
        EL_ERROR("libubi_open failed in AttachMtd");
        return false;
    }
    struct ubi_attach_request req{};
    req.dev_num = UBI_DEV_NUM_AUTO;
    req.mtd_num = mtd_num;
    req.mtd_dev_node = nullptr;
    req.vid_hdr_offset = 0;
    req.max_beb_per1024 = 0; // default
    int rc = ubi_attach(ubi, kUbiCtrl, &req);
    if (rc < 0) {
        EL_ERROR("ubi_attach(mtd{}): {}", mtd_num, strerror(errno));
        libubi_close(ubi);
        return false;
    }
    ubi_num = req.dev_num;
    libubi_close(ubi);
    return true;
}

bool UbiManager::GetVolIdByName(int ubi_num, const std::string &vol_name, int &vol_id)
{
    libubi_t ubi = libubi_open();
    if (!ubi) return false;
    struct ubi_vol_info vinfo{};
    int rc = ubi_get_vol_info1_nm(ubi, ubi_num, vol_name.c_str(), &vinfo);
    libubi_close(ubi);
    if (rc != 0) return false;
    vol_id = vinfo.vol_id;
    return true;
}

bool UbiManager::CreateUbiblockAndMount(int ubi_num, int vol_id, const std::string &mount_point)
{
    char vol_node[64];
    snprintf(vol_node, sizeof(vol_node), "/dev/ubi%d_%d", ubi_num, vol_id);

    int fd = open(vol_node, O_RDONLY);
    if (fd < 0) {
        EL_ERROR("open {} failed: {}", vol_node, strerror(errno));
        return false;
    }

    if (ubi_vol_block_create(fd) != 0) {
        // it may already exist; continue
        EL_WARN("ubi_vol_block_create {}: {} (may already exist)", vol_node, strerror(errno));
    }
    close(fd);

    char blk_node[64];
    snprintf(blk_node, sizeof(blk_node), "/dev/ubiblock%d_%d", ubi_num, vol_id);

    // Ensure mount point exists
    ::mkdir(mount_point.c_str(), 0755);
    if (mount(blk_node, mount_point.c_str(), "squashfs", MS_RDONLY | MS_NOSUID | MS_NODEV, nullptr) != 0) {
        EL_ERROR("mount {} -> {} failed: {}", blk_node, mount_point, strerror(errno));
        return false;
    }
    return true;
}

bool UbiManager::PrepareStaticVolumePreWrite(int mtd_num, const std::string &mount_point, const std::string &vol_name)
{
    (void)vol_name; // Not needed for pre-write now (we assume vol_id=0)

    // 1) Unmount if mounted
    if (!UmountIfMounted(mount_point)) {
        return false;
    }

    // 2) If attached, remove ubiblock and detach
    int ubi_num = -1;
    if (FindUbiNumByMtd(mtd_num, ubi_num)) {
        // remove ubiblock for vol 0 (static volume created by packaging)
        (void)RemoveUbiblockIfExists(ubi_num, 0);
        (void)DetachMtd(mtd_num);
    }
    return true;
}

bool UbiManager::RestoreStaticVolumePostWrite(int mtd_num, const std::string &mount_point, const std::string &vol_name)
{
    int ubi_num = -1;
    if (!AttachMtd(mtd_num, ubi_num)) {
        EL_ERROR("Failed to re-attach mtd{} to UBI", mtd_num);
        return false;
    }

    int vol_id = 0;
    // Prefer lookup by name; fallback to 0
    if (!GetVolIdByName(ubi_num, vol_name, vol_id)) {
        EL_WARN("Cannot find volume '{}' on ubi{}, fallback to vol_id=0", vol_name, ubi_num);
        vol_id = 0;
    }

    if (!CreateUbiblockAndMount(ubi_num, vol_id, mount_point)) {
        EL_ERROR("Failed to create ubiblock/mount for ubi{}_{} -> {}", ubi_num, vol_id, mount_point);
        return false;
    }
    return true;
}

} // namespace Upgrade
} // namespace El

